#include "Cholesky.h"
#include "Vector.h"
#include <cmath>
#include <iomanip>

using namespace RSIM;

#ifndef RSIM_SMALL
        #define RSIM_SMALL 1.0E-6
#endif

///////////////////////////////////////////////////////	

Cholesky::Cholesky(const Matrix &A){
        const int& m = A.dim1();
        const int& n = A.dim2();
        
        #ifdef RSIM_SAFE
        isspd = (m == n);
        
        if (m != n){
                cout<<"Warning:Input Matrix for the constructor Cholesky(const Matrix& A) is not square.\n";
                return;
        }
        #endif
        
        L_.resize(n,n);
        double *lptr = L_.wPtr();
        const double *aptr  = A.Ptr();
        
        // Main loop.
        for(int j=0; j<n; j++){
                double d = 0.0;
                for(int k=0; k<j; k++){
                        double s= 0.0;
                        for (int i=0; i<k; i++){
                               // s += L_[k][i]*L_[j][i];
                               s += lptr[k*n+i]*lptr[j*n+i];
                        }
                        // L_[j][k] = s = (A[j][k] - s)/L_[k][k];                       
                       s = (aptr[j*n+k] - s)/lptr[k*n+k];
                       lptr[j*n+k] = s;
                        d = d + s*s;
                        // isspd = isspd && (A[k][j] == A[j][k]);
                       isspd = isspd && (fabs(aptr[k*n+j]-aptr[j*n+k])<RSIM_SMALL); 
                }                
                // d = A[j][j] - d;
                d = aptr[j*n+j] - d;
                isspd = isspd && (d > 0.0);
                // L_[j][j] = sqrt(d > 0.0 ? d : 0.0);
                lptr[j*n+j] = sqrt(d > 0.0 ? d : 0.0);
                for(int k = j+1; k < n; k++){
                        // L_[j][k] = 0.0;
                        lptr[j*n+k] = 0.0;
                }
        }        
}

///////////////////////////////////////////////////////	

Vector Cholesky::solve(const Vector& b){
        int n = L_.dim1();
        if (b.getLength() != n)
                return Vector();
        
        Vector x = b;
        
        // Solve L*y = b;
        for(int k = 0; k < n; k++){
                for(int i = 0; i < k; i++) 
                        x[k] -= x[i]*L_[k][i];
                x[k] /= L_[k][k];
        }
        
        // Solve L'*X = Y;
        for(int k = n-1; k >= 0; k--){
                for(int i = k+1; i < n; i++) 
                        x[k] -= x[i]*L_[i][k];
                x[k] /= L_[k][k];
        }
        return x;
}

///////////////////////////////////////////////////////	

void Cholesky::solve(const Matrix& B, Matrix& X){
                
        int n = L_.dim1();
        #ifdef RSIM_SAFE
        if (B.dim1() != n){
                cout<<"Warning: Check dimensions of input matrices to the object of class Cholesky.\n";
        }
        #endif
        
        X.resize(B.dim1(),B.dim2());
        double *x = X.wPtr();
        const double *bptr = B.Ptr();

        for(int i=0;i<B.dim1()*B.dim2();++i)
                x[i]=bptr[i];

        int nx = B.dim2();
        // Solve L*y = b;
        for(int j=0; j< nx; j++){
                for(int k = 0; k < n; k++) 
                {
                        for(int i = 0; i < k; i++) 
                                X[k][j] -= X[i][j]*L_[k][i];
                        X[k][j] /= L_[k][k];
                }
        }
        
        // Solve L'*X = Y;
        for(int j=0; j<nx; j++){
                for(int k = n-1; k >= 0; k--){
                        for (int i = k+1; i < n; i++) 
                                X[k][j] -= X[i][j]*L_[i][k];
                        X[k][j] /= L_[k][k];
                }
        }
}

///////////////////////////////////////////////////////	

void Cholesky::solve(const Matrix& B, double *X){
        
        const int& n = L_.dim1();                        
        const double *bptr = B.Ptr();
        const double *lptr = L_.Ptr();

        for(int i=0;i<n*n;++i)
                X[i]=bptr[i];

        int nx = B.dim2();
        // Solve L*y = b;
        for(int j=0; j< nx; j++){
                for(int k = 0; k < n; k++) 
                {
                        for(int i = 0; i < k; i++) 
                                X[k*n+j] -= X[i*n+j]*lptr[k*n+i];
                        X[k*n+j] /= lptr[k*n+k];
                }
        }
        
        // Solve L'*X = Y;
        for(int j=0; j<nx; j++){
                for(int k = n-1; k >= 0; k--){
                        for (int i = k+1; i < n; i++) 
                                X[k*n+j] -= X[i*n+j]*lptr[i*n+k];
                        X[k*n+j] /= lptr[k*n+k];
                }
        }
}

void Cholesky::inverse(double *X){
 
        #ifdef RSIM_SAFE        
        assert(isspd);
        #endif
       
        const int& n = L_.dim1();                                
        const double *lptr = L_.Ptr();

        for(int i=0;i<n*n;X[i++]=0.0);
        for(int i=0;i<n;++i)
                X[n*i+i]=1.0;
        
        // Solve L*y = b;
        for(int j=0;j<n;j++){                
                for(int k = 0; k < n; k++) {
                        for(int i = 0; i < k; i++){ 
                                X[k*n+j] -= X[i*n+j]*lptr[k*n+i];
                        }
                        X[k*n+j] /= lptr[k*n+k];
                }
        }
                        
        // Solve L'*X = Y;
        for(int j=0; j<n; j++){
                for(int k = n-1; k >= 0; k--){
                        for (int i = k+1; i < n; i++) 
                                X[k*n+j] -= X[i*n+j]*lptr[i*n+k];
                        X[k*n+j] /= lptr[k*n+k];
                }
        }
}